import entity.HuffCodeNode;

import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * @Author: wanqian
 * @Date: 2021/2/8 17:02
 */
public class HuffmanEncodeDemo {

    public static String content = "i like like like java do you like a java";

    public static void main(String[] args) {
        //将字符串转换成赫夫曼编码的二级制数，以字节数组保存
        //将最初40个字符(320bit)编码成17个字节的二进制数(136bit)，压缩率42.5%
        byte[] bytes = huffmanZip(content.getBytes());
        System.out.println(Arrays.toString(bytes));
    }

    public static byte[] huffmanZip(byte[] bytes) {
        //将字符串转换为赫夫曼树
        HuffCodeNode rootNode = huffmanTree(bytes);
        //获得对应的赫夫曼编码表
        Map<Byte, String> huffCodeTable = getCodes(rootNode);

        StringBuilder stringBuilder = new StringBuilder();
        for (byte b : bytes) {
            String code = huffCodeTable.get(b);
            stringBuilder.append(code);
        }
        //10101000 -> 10100111 -> 11011000
        //1. 得到赫夫曼字符串 1010100010111111...
        //2. 此时二进制数都是以字符保存(8bit)，直接传输时显然不合理
        //3. 将其转换成bit存储一个二进制数，即八个字符转成一个byte保存，最终为byte[]进行传输
        String huffCodeString = stringBuilder.toString();

        //byte[]的长度：八位字符转为一个字节，末尾不足8的位补1
        int stringLength = huffCodeString.length();
        int length = (stringLength + 7) / 8;

        byte[] huffBytes = new byte[length + 1];
        int index = 0;
        int bitLength = 0;
        for (int i = 0; i < stringLength; i += 8) {
            String codeString;
            //取八个字符，末尾不足的取到最后
            bitLength = stringLength - i;
            if (bitLength >= 8) {
                codeString = huffCodeString.substring(i, i + 8);
            } else {
                codeString = huffCodeString.substring(i);
            }
            //将字符以2进制转成int，然后强转成byte，进行保存
            huffBytes[index++] = (byte) Integer.parseInt(codeString, 2);
        }
        huffBytes[length] = (byte) bitLength;
        return huffBytes;
    }

    public static Map<Byte, String> getHuffCodeTable(String content) {
        HuffCodeNode root = huffmanTree(content);
        return getCodes(root);
    }

    public static Map<Byte, String> getHuffCodeTable(byte[] content) {
        HuffCodeNode root = huffmanTree(content);
        return getCodes(root);
    }

    /**
     * 根据字符串创建赫夫曼树
     *
     * @param content
     * @return
     */
    public static HuffCodeNode huffmanTree(String content) {

        //统计出现的个数
        Map<Integer, Long> collect = content.chars().mapToObj(Integer::new).collect(
                Collectors.groupingBy(Function.identity(), Collectors.counting())
        );
        //创建节点
        List<HuffCodeNode> nodes = collect.entrySet().stream()
                .map(entry -> new HuffCodeNode(entry.getKey().byteValue(), entry.getValue().intValue()))
                .collect(Collectors.toList());
        //创建赫夫曼树
        while (nodes.size() > 1) {
            Collections.sort(nodes);
            //取出最小的两个节点
            HuffCodeNode leftNode = nodes.remove(0);
            HuffCodeNode rightNode = nodes.remove(0);
            //生成一个父节点并放入
            HuffCodeNode parent = new HuffCodeNode(leftNode, rightNode);
            nodes.add(parent);
        }

        return nodes.get(0);
    }

    /**
     * 根据字符串创建赫夫曼树
     *
     * @param bytes
     * @return
     */
    public static HuffCodeNode huffmanTree(byte[] bytes) {

        ArrayList<Byte> characters = new ArrayList<>();
        for (byte b : bytes) {
            characters.add(b);
        }
        //统计出现的个数
        Map<Byte, Long> collect = characters.stream().collect(
                Collectors.groupingBy(Function.identity(), Collectors.counting())
        );
        //创建节点
        List<HuffCodeNode> nodes = collect.entrySet().stream()
                .map(entry -> new HuffCodeNode(entry.getKey().byteValue(), entry.getValue().intValue()))
                .collect(Collectors.toList());
        //创建赫夫曼树
        while (nodes.size() > 1) {
            Collections.sort(nodes);
            //取出最小的两个节点
            HuffCodeNode leftNode = nodes.remove(0);
            HuffCodeNode rightNode = nodes.remove(0);
            //生成一个父节点并放入
            HuffCodeNode parent = new HuffCodeNode(leftNode, rightNode);
            nodes.add(parent);
        }

        return nodes.get(0);
    }


    /**
     * 获取赫夫曼编码表
     *
     * @param root
     * @return
     */
    public static Map<Byte, String> getCodes(HuffCodeNode root) {
        if (root == null) {
            return null;
        }
        HashMap<Byte, String> huffCodeTable = new HashMap<>();
        getCode(root.getLeft(), "0", new StringBuilder(), huffCodeTable);
        getCode(root.getRight(), "1", new StringBuilder(), huffCodeTable);
        return huffCodeTable;
    }

    /**
     * 获得每个字符对应的赫夫曼编码
     *
     * @param node
     * @param pathCode
     * @return
     */
    public static void getCode(HuffCodeNode node, String pathCode, StringBuilder stringBuilder, HashMap<Byte, String> huffCodeTable) {
        if (node != null) {
            // 此处是关键
            // 既要继承节点之前的路径码又要保证两个路径不互相影响
            StringBuilder sb = new StringBuilder(stringBuilder);
            sb.append(pathCode);
            if (node.getData() != null) {
                huffCodeTable.put(node.getData(), sb.toString());
            } else {
                getCode(node.getLeft(), "0", sb, huffCodeTable);
                getCode(node.getRight(), "1", sb, huffCodeTable);
            }
        }
    }
}
